package com.example.videopractice.mediaex

import android.media.MediaCodec
import android.media.MediaExtractor
import android.media.MediaFormat
import android.media.MediaMuxer
import android.os.Environment
import com.example.base.util.ALog
import java.io.File
import java.nio.ByteBuffer

/**
 * 将录制好的mp4音频和视频混合成一个视频
 */
class H264Muxer(
    private val videoPath: String = Environment.getExternalStorageDirectory().absolutePath + File.separator + "temp_video.mp4",
    private val audioPath: String = Environment.getExternalStorageDirectory().absolutePath + File.separator + "temp_audio.mp4"
) {

    private val outputPath =
        Environment.getExternalStorageDirectory().absolutePath + File.separator + "video_output.mp4"
    private val videoExtractor = MediaExtractor()
    private val audioExtractor = MediaExtractor()
    private var mediaMuxer:MediaMuxer? = null

    // 开始混合
    @Throws(RuntimeException::class)
    fun muxer() {
        mediaMuxer = MediaMuxer(outputPath, MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4)
        val muxer = mediaMuxer ?: return

        //音轨信息
        var audioTrackIndex = -1
        var audioMuxerTrackIndex = -1
        var audioMaxInputSize = 0
        audioExtractor.setDataSource(audioPath)
        val audioTrackCount = audioExtractor.trackCount
        for (i in 0 until audioTrackCount) {
            val format = audioExtractor.getTrackFormat(i)
            val mime = format.getString(MediaFormat.KEY_MIME) ?: continue
            if (mime.startsWith("audio/")) {
                ALog.e("xiao","找到音频轨")
                audioTrackIndex = i
                audioMuxerTrackIndex = muxer.addTrack(format)//将音轨添加到MediaMuxer，并返回新的轨道
                audioMaxInputSize = format.getInteger(MediaFormat.KEY_MAX_INPUT_SIZE)
                break
            }
        }

        //图像信息
        var videoTrackIndex = -1
        var videoMuxerTrackIndex = -1
        var videoMaxInputSize = 0
        var videoFrameRate = 0
        videoExtractor.setDataSource(videoPath)
        val videoTrackCount = videoExtractor.trackCount
        for (i in 0 until videoTrackCount) {
            val format = videoExtractor.getTrackFormat(i)
            val mime = format.getString(MediaFormat.KEY_MIME) ?: continue
            if (mime.startsWith("video/")) {
                ALog.e("xiao","找到视频轨")
                videoTrackIndex = i
                videoMuxerTrackIndex = muxer.addTrack(format) //将视频轨添加到MediaMuxer，并返回新的轨道
                videoMaxInputSize = format.getInteger(MediaFormat.KEY_MAX_INPUT_SIZE); //得到能获取的有关视频的最大值
                videoFrameRate = format.getInteger(MediaFormat.KEY_FRAME_RATE)//获取视频的帧率
                break
            }
        }

        if (audioTrackIndex == -1) throw RuntimeException("未找到音轨")
        if (videoTrackIndex == -1) throw RuntimeException("未找到视频轨")

        muxer.start() //开始合成

        audioExtractor.selectTrack(audioTrackIndex) //将提供音频的视频选择到音轨上
        val audioMediaInfo = MediaCodec.BufferInfo()
        val audioBuffer = ByteBuffer.allocate(audioMaxInputSize)
        while (true) {
            val sampleSize = audioExtractor.readSampleData(audioBuffer, 0)//检索当前编码的样本并将其存储在字节缓冲区中
            if (sampleSize <= 0) {
                audioExtractor.unselectTrack(audioTrackIndex)
                break
            }

            //设置样本编码信息
            audioMediaInfo.offset = 0
            audioMediaInfo.size = sampleSize
            audioMediaInfo.flags = audioExtractor.sampleFlags
            audioMediaInfo.presentationTimeUs = audioExtractor.sampleTime

            muxer.writeSampleData(audioMuxerTrackIndex, audioBuffer, audioMediaInfo)
            audioExtractor.advance()
        }

        videoExtractor.selectTrack(videoTrackIndex) ////将提供视频图像的视频选择到视频轨上
        val videoMediaInfo = MediaCodec.BufferInfo()
        val videoBuffer = ByteBuffer.allocate(videoMaxInputSize)
        while (true) {
            val sampleSize = videoExtractor.readSampleData(videoBuffer, 0)//检索当前编码的样本并将其存储在字节缓冲区中
            if (sampleSize <= 0) {
                videoExtractor.unselectTrack(videoTrackIndex)
                break
            }

            //设置样本编码信息
            videoMediaInfo.offset = 0
            videoMediaInfo.size = sampleSize
            videoMediaInfo.flags = MediaCodec.BUFFER_FLAG_KEY_FRAME
            videoMediaInfo.presentationTimeUs += 1000 * 1000 / videoFrameRate

            muxer.writeSampleData(videoMuxerTrackIndex, videoBuffer, videoMediaInfo)
            videoExtractor.advance()
        }

        audioExtractor.release()
        videoExtractor.release()

        muxer.stop()
        muxer.release()
        ALog.e("xiao","完成了")
    }
}